package simpletrace

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"unsafe"
)

func GetFuncName(f any) string {
	fn := reflect.ValueOf(f).Pointer()
	fullName := runtime.FuncForPC(fn).Name() // e.g. "main.RealAdd"
	if i := strings.LastIndex(fullName, "."); i != -1 {
		return strings.ReplaceAll(fullName[i+1:], "-fm", "") // 取最后一段名称
	}
	return strings.ReplaceAll(fullName, "-fm", "")
}

func shouldReturnFrame(frame runtime.Frame, fn string) bool {
	// 特殊情况：总是返回这些特定路径
	if strings.Contains(frame.File, "simpletrace/util_test") ||
		strings.Contains(frame.File, "simpletrace/cache") {
		return true
	}

	// 排除的文件路径
	excludedPaths := []string{
		"simpletrace", "github", "pkg/", "chat/lib", "go/src",
	}

	// 排除的函数特征
	excludedFunctions := []string{
		"Trace", "SafeGo",
	}

	// 检查文件路径
	containsExcludedPath := false
	for _, path := range excludedPaths {
		if strings.Contains(frame.File, path) {
			containsExcludedPath = true
			break
		}
	}

	// 检查函数名称
	containsExcludedFunc := false
	for _, funcName := range excludedFunctions {
		if strings.Contains(frame.Function, funcName) {
			containsExcludedFunc = true
			break
		}
	}

	return !containsExcludedPath && !containsExcludedFunc && !strings.HasPrefix(fn, "func")
}

func getCallerFileLine() (string, string, int) {
	pc := make([]uintptr, 20)   // 最多 20 层栈
	n := runtime.Callers(2, pc) // skip runtime.Callers + getRealCallerFileLine
	frames := runtime.CallersFrames(pc[:n])
	var lastFrame runtime.Frame
	var lastFn string

	for {
		frame, more := frames.Next()
		if frame.Function == "" {
			break
		}
		lastFrame = frame
		fns := strings.Split(frame.Function, ".")
		fn := fns[len(fns)-1]
		lastFn = fn
		if shouldReturnFrame(frame, fn) {
			return fn, frame.File, frame.Line
		}
		if !more {
			break
		}
	}
	return lastFn, lastFrame.File, lastFrame.Line
}

func Trace[T any](f T) T {
	return WrapFunc(f)
}

func WrapFunc[T any](f T) T {
	return Wrap(f).(T)
}

func WrapFuncNoCache[T any](f T) T {
	return Wrap(f, &WrapOption{WithoutCache: true}).(T)
}

func WrapFuncWithCacheKey[T any](key string, f T) T {
	return Wrap(f, &WrapOption{CacheKey: key}).(T)
}

func TraceStart[T any](f T) T {
	return WrapFuncStart(f)
}

func WrapFuncStart[T any](f T) T {
	return WrapStart(f).(T)
}

func WrapFuncNoCacheStart[T any](f T) T {
	return WrapStart(f, &WrapOption{WithoutCache: true}).(T)
}

func WrapFuncWithCacheKeyStart[T any](key string, f T) T {
	return WrapStart(f, &WrapOption{CacheKey: key}).(T)
}

type WrapStartFn func(ctx context.Context, name string, pos string, args map[string]interface{}) (context.Context, WrapEndFn)
type WrapEndFn func(ctx context.Context, returns map[string]interface{}, hasError ...bool)

func DefaultWrapStartFn(ctx context.Context, name string, pos string, args map[string]interface{}) (context.Context, WrapEndFn) {
	span, ctx := StartSpan(ctx, name, args, true, WithPos(pos))
	return ctx, func(ctx context.Context, returns map[string]interface{}, hasError ...bool) {
		span.Finish(ctx, returns, hasError...)
	}
}

type WrapOption struct {
	WithoutCache bool   // 是否不使用缓存
	CacheKey     string // 缓存的键

	StartEndFn WrapStartFn
}

func GetParentIDAndSpanIDFromCtx(ctx context.Context) (string, string, string) {
	traceID, ok := ctx.Value(CtxHeaderSimpleTraceID).(string)
	if !ok || traceID == "" {
		return "", "", ""
	}

	parentID, ok := ctx.Value(CtxHeaderSimpleTraceParentID).(string)
	if !ok || parentID == "" {
		parentID = ""
	}

	spanID, ok := ctx.Value(CtxHeaderSimpleTraceSpanID).(string)
	if !ok || spanID == "" {
		spanID = ""
	}

	return traceID, parentID, spanID
}

var cachedWrapFunc sync.Map

func WrapStart(f any, opt ...*WrapOption) any {
	wrapOpt := &WrapOption{}
	if len(opt) > 0 {
		wrapOpt = opt[0]
	}
	if wrapOpt.StartEndFn == nil {
		wrapOpt.StartEndFn = func(ctx context.Context, name, pos string, args map[string]interface{}) (context.Context, WrapEndFn) {
			traceID, parentID, spanID := GetParentIDAndSpanIDFromCtx(ctx)
			span, ctx := StartSpan(ctx, name, args, true, WithIDs(traceID, spanID, parentID))
			return ctx, func(ctx context.Context, returns map[string]interface{}, hasError ...bool) {
				span.Finish(ctx, returns, hasError...)
			}
		}
	}
	return wrap(f, wrapOpt)
}

func Wrap(f any, opt ...*WrapOption) any {
	return wrap(f, opt...)
}

func wrap(f any, opt ...*WrapOption) any {
	fnValue := reflect.ValueOf(f)
	var ptr = fnValue.Pointer()

	var withoutCache bool
	var cacheKey string
	if len(opt) > 0 {
		withoutCache = opt[0].WithoutCache
		cacheKey = opt[0].CacheKey
	}

	if !withoutCache {
		if cacheKey != "" {
			if cached, ok := cachedWrapFunc.Load(cacheKey); ok {
				return cached
			}
		} else {
			if cached, ok := cachedWrapFunc.Load(ptr); ok {
				return cached
			}
		}
	}

	fnType := fnValue.Type()
	// 如果不是函数类型，直接返回原值
	if fnType.Kind() != reflect.Func {
		return f
	}

	// 第一个参数不是 context.Context 或 *gin.Context 就直接返回原函数
	if fnType.NumIn() == 0 || (fnType.In(0).String() != "context.Context" && fnType.In(0).String() != "*gin.Context") {
		return f
	}

	fnName := GetFuncName(f)

	// 防止嵌套调用
	if excludeWrapFunc[strings.ToLower(fnName)] {
		return f
	}

	caller, file, line := getCallerFileLine()

	if caller == "ReStoreTraceContext" {
		return f
	}

	fnPos := fmt.Sprintf("%s:%d", file, line)

	fnName = getWrapFuncName(fnName, caller)

	wrapped := reflect.MakeFunc(fnType, func(args []reflect.Value) (results []reflect.Value) {
		var ctx context.Context
		if c, ok := args[0].Interface().(context.Context); ok {
			ctx = c
		}

		tobecallFn := callFunc

		argMap := map[string]interface{}{}
		for i := 1; i < len(args); i++ {
			if args[i].Type().String() == "[]uint8" || args[i].Type().String() == "[]byte" {
				argMap[fmt.Sprintf("arg%d", i)] = string(args[i].Bytes())
				continue
			}
			// 判断是否为不可导出和json不可序列化的类型，比如函数等，是则跳过
			if args[i].Type().Kind() == reflect.Func || args[i].Type().Kind() == reflect.Chan ||
				(args[i].Type().Kind() == reflect.Ptr && args[i].Type().Elem().Kind() == reflect.Func) {
				argMap[fmt.Sprintf("arg%d", i)] = "<func>"
				continue
			}
			if args[i].Type().Kind() == reflect.Ptr && args[i].IsNil() {
				argMap[fmt.Sprintf("arg%d", i)] = "<nil>"
				continue
			}

			argMap[fmt.Sprintf("arg%d", i)] = args[i].Interface()
		}

		startFn := DefaultWrapStartFn
		if len(opt) > 0 && opt[0].StartEndFn != nil {
			startFn = opt[0].StartEndFn
		}

		ctx, endFn := startFn(ctx, fnName, fnPos, argMap)
		// 如果第一个参数是 context.Context， 则进行替换
		if fnType.In(0).String() == "context.Context" {
			args[0] = reflect.ValueOf(ctx)
		}

		results = tobecallFn(fnValue, fnType, args)

		hasError := false
		retMap := map[string]interface{}{}
		for i := range results {
			// 如果是ctx类型，则跳过
			if results[i].Type().String() == "context.Context" || results[i].Type().String() == "*gin.Context" {
				retMap[fmt.Sprintf("ret%d", i)] = "<ctx>"
				continue
			}
			retMap[fmt.Sprintf("ret%d", i)] = results[i].Interface()

			if i == len(results)-1 {
				// 如果最后一个返回值是error类型，则判断是否有错误
				if results[i].Type().String() == "error" {
					if err, ok := results[i].Interface().(error); ok && err != nil {
						hasError = true
					}
				}
			}
		}
		endFn(ctx, retMap, hasError)
		return results
	})

	result := wrapped.Interface()

	if !withoutCache {
		if cacheKey != "" {
			cachedWrapFunc.Store(cacheKey, result)
			return result
		}
		cachedWrapFunc.Store(ptr, result)
	}

	return result
}

func callFunc(fnValue reflect.Value, fnType reflect.Type, args []reflect.Value) []reflect.Value {
	if fnType.IsVariadic() {
		numFixed := fnType.NumIn() - 1
		fixedArgs := args[:numFixed]
		variadicArgs := args[numFixed]

		for i := range variadicArgs.Len() {
			fixedArgs = append(fixedArgs, variadicArgs.Index(i))
		}
		return fnValue.Call(fixedArgs)
	}
	return fnValue.Call(args)
}

func GetStructSize(v any) string {
	size := Sizeof(v)
	return FormatSize(size)
}

func FormatMillseconds(ms int64) string {
	if ms < 1000 {
		return fmt.Sprintf("%d ms", ms)
	}
	seconds := float64(ms) / 1000
	if seconds < 60 {
		return fmt.Sprintf("%.2f s", seconds)
	}
	minutes := seconds / 60
	if minutes < 60 {
		return fmt.Sprintf("%.2f min", minutes)
	}
	hours := minutes / 60
	return fmt.Sprintf("%.2f h", hours)
}

func FormatSize(size uintptr) string {
	const (
		_          = iota
		KB float64 = 1 << (10 * iota)
		MB
		GB
	)

	s := float64(size)
	switch {
	case s >= GB:
		return fmt.Sprintf("%.2f GB", s/GB)
	case s >= MB:
		return fmt.Sprintf("%.2f MB", s/MB)
	case s >= KB:
		return fmt.Sprintf("%.2f KB", s/KB)
	default:
		return fmt.Sprintf("%d B", size)
	}
}

// visited 用于避免循环引用导致无限递归
type visited map[uintptr]struct{}

// Sizeof 用于计算结构体占用的内存（递归版本）
func Sizeof(v any) uintptr {
	vis := visited{}
	return sizeofRecursive(reflect.ValueOf(v), vis)
}

func sizeofRecursive(v reflect.Value, vis visited) uintptr {
	if !v.IsValid() {
		return 0
	}

	// 解引用指针
	for v.Kind() == reflect.Ptr {
		ptr := v.Pointer()
		if ptr == 0 || vis[ptr] != (struct{}{}) {
			return 0
		}
		vis[ptr] = struct{}{}
		v = v.Elem()
	}

	switch v.Kind() {
	case reflect.Struct:
		var size uintptr
		for i := 0; i < v.NumField(); i++ {
			size += sizeofRecursive(v.Field(i), vis)
		}
		return size
	case reflect.Slice:
		var size uintptr = unsafe.Sizeof(v.Interface())
		for i := 0; i < v.Len(); i++ {
			size += sizeofRecursive(v.Index(i), vis)
		}
		return size
	case reflect.Map:
		var size uintptr = unsafe.Sizeof(v.Interface())
		for _, key := range v.MapKeys() {
			size += sizeofRecursive(key, vis)
			size += sizeofRecursive(v.MapIndex(key), vis)
		}
		return size
	case reflect.String:
		return unsafe.Sizeof(v.Interface()) + uintptr(len(v.String()))
	default:
		return unsafe.Sizeof(v.Interface())
	}
}

func getWrapFuncName(funcName, caller string) string {
	if !needWrapFuncName[funcName] {
		return funcName
	}
	return fmt.Sprintf("%s.%s", caller, funcName)
}

var needWrapFuncName = map[string]bool{
	"GetWithCtx":                           true,
	"SetWithCtx":                           true,
	"GetSetWithCtx":                        true,
	"SetExWithCtx":                         true,
	"SetNxWithCtx":                         true,
	"SetNxExWithCtx":                       true,
	"MGetWithCtx":                          true,
	"MSetWithCtx":                          true,
	"MSetNXWithCtx":                        true,
	"SetBitWithCtx":                        true,
	"DelWithCtx":                           true,
	"IncrWithCtx":                          true,
	"IncrByWithCtx":                        true,
	"ExpireWithCtx":                        true,
	"ExpireAtWithCtx":                      true,
	"TTlWithCtx":                           true,
	"ExistsWithCtx":                        true,
	"SetBitV2WithCtx":                      true,
	"GetBitWithCtx":                        true,
	"BitCountWithCtx":                      true,
	"BitOpAndWithCtx":                      true,
	"BitOpOrWithCtx":                       true,
	"BitOpXorWithCtx":                      true,
	"BitOpNotWithCtx":                      true,
	"ScanWithCtx":                          true,
	"HGetWithCtx":                          true,
	"HSetWithCtx":                          true,
	"HMGetWithCtx":                         true,
	"HMSetWithCtx":                         true,
	"HExpireWithCtx":                       true,
	"HKeysWithCtx":                         true,
	"HValsWithCtx":                         true,
	"HGetAllWithCtx":                       true,
	"HIncrByWithCtx":                       true,
	"HDelWithCtx":                          true,
	"HClearWithCtx":                        true,
	"HLenWithCtx":                          true,
	"SAddWithCtx":                          true,
	"SPopWithCtx":                          true,
	"SIsMemberWithCtx":                     true,
	"SRemWithCtx":                          true,
	"SMembersWithCtx":                      true,
	"SCardWithCtx":                         true,
	"SScanWithCtx":                         true,
	"ZAddWithCtx":                          true,
	"ZAddNXWithCtx":                        true,
	"ZAddXXWithCtx":                        true,
	"ZRemWithCtx":                          true,
	"ZRemRangeByScoreWithCtx":              true,
	"ZExpireWithCtx":                       true,
	"ZRangeWithCtx":                        true,
	"ZRevRangeWithCtx":                     true,
	"ZRangeByScoreWithScoreWithCtx":        true,
	"ZRangeByScoreWithScoreLimitedWithCtx": true,
	"ZRevRangeByScoreWithScoreWithCtx":     true,
	"ZClearWithCtx":                        true,
	"ZCountWithCtx":                        true,
	"ZCardWithCtx":                         true,
	"ZScoreWithCtx":                        true,
	"ZRankWithCtx":                         true,
	"ZRevRankWithCtx":                      true,
	"LRangeWithCtx":                        true,
	"LLenWithCtx":                          true,
	"LRemWithCtx":                          true,
	"LPopWithCtx":                          true,
	"RPopWithCtx":                          true,
	"BLPopWithCtx":                         true,
	"BRPopWithCtx":                         true,
	"LPushWithCtx":                         true,
	"RPushWithCtx":                         true,
	"BRPopLPushWithCtx":                    true,
	"RPopLPushWithCtx":                     true,
	"LSetWithCtx":                          true,
	"ScriptLoadWithCtx":                    true,
	"EvalShaWithCtx":                       true,
	"EvalWithCtx":                          true,
	"PFAddWithCtx":                         true,
	"PFCountWithCtx":                       true,
	"GetInt64WithCtx":                      true,
	"MGetWithPipelineWithCtx":              true,
	"QueryWithCtx":                         true,
	"ExecWithCtx":                          true,
	"QueryInt64WithCtx":                    true,
	"QueryStringWithCtx":                   true,
}

var excludeWrapFunc = map[string]bool{
	"trace":                true,
	"wrap":                 true,
	"wrapfunc":             true,
	"wrapfuncnocache":      true,
	"tracestart":           true,
	"wrapstart":            true,
	"wrapfuncstart":        true,
	"wrapfuncnocachestart": true,
	"wrapfuncwithcachekey": true,
}

func SafeMarshal(v interface{}) ([]byte, error) {
	cleaned := cleanUnsupported(reflect.ValueOf(v))
	return json.Marshal(cleaned)
}

func cleanUnsupported(v reflect.Value) interface{} {
	switch v.Kind() {
	case reflect.Struct:
		m := make(map[string]interface{})
		t := v.Type()
		for i := 0; i < v.NumField(); i++ {
			field := t.Field(i)
			if field.PkgPath != "" { // 非导出字段跳过
				continue
			}
			val := v.Field(i)
			cleanedVal := cleanUnsupported(val)
			if cleanedVal != nil {
				jsonKey := field.Name
				if tag := field.Tag.Get("json"); tag != "" && tag != "-" {
					jsonKey = tag
				}
				m[jsonKey] = cleanedVal
			}
		}
		return m
	case reflect.Map:
		m := make(map[string]interface{})
		for _, key := range v.MapKeys() {
			m[key.String()] = cleanUnsupported(v.MapIndex(key))
		}
		return m
	case reflect.Slice, reflect.Array:
		arr := make([]interface{}, v.Len())
		for i := 0; i < v.Len(); i++ {
			arr[i] = cleanUnsupported(v.Index(i))
		}
		return arr
	default:
		if isSupportedKind(v.Kind()) {
			return v.Interface()
		}
		return nil // 不支持的直接忽略
	}
}

func isSupportedKind(k reflect.Kind) bool {
	switch k {
	case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16,
		reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8,
		reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr,
		reflect.Float32, reflect.Float64, reflect.String:
		return true
	default:
		return false
	}
}
